use std::fs::File;
use std::io::{Read, Write};
use std::os::unix::net::{UnixListener, UnixStream};
use std::path::PathBuf;

use clap::Parser;
use log::error;

#[derive(Parser)]
enum Cli {
    /// Start service
    StartSrv {
        /// Remove socket file before start.
        #[clap(short, long)]
        force: bool,
    },
    /// Stop the service.
    StopSrv,
    /// Add a task to the running service.
    Add {
        prog: String,
        args: Vec<String>,

        /// Set a name to the program.
        #[clap(short, long)]
        name: Option<String>,
    },
    /// Show status of a task.
    Stat {
        /// Task ID.
        id: usize,
    },
    /// Show output of task with ID.
    Log {
        /// Task ID.
        id: usize,

        /// Show STDERR instead of STDOUT.
        #[clap(short = 'e', long)]
        stderr: bool,
    },
    /// Kill a child proccess.
    Kill {
        /// Task ID.
        id: usize,
    },
}

#[derive(Debug, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
enum Status {
    Running,
    Done(i32),
    Killed,
}

#[derive(Debug, serde::Serialize, serde::Deserialize)]
struct TaskInfo {
    show_name: Option<String>,
    pid: u32,

    id: usize,
    status: Status,
    prog: String,
    args: Vec<String>,
    cwd: PathBuf,
    envs: Vec<(String, String)>,
}

impl TaskInfo {
    fn change_status(&mut self, status: Status) -> Result<(), ()> {
        if status == self.status {
            return Ok(());
        }
        self.status = status;
        dump_task_info(&self)?;
        Ok(())
    }
}

#[derive(Debug)]
struct Task {
    info: TaskInfo,
    child: std::process::Child,
}

fn main() {
    if runtime().is_err() {
        std::process::exit(1);
    }
}

fn data_dir() -> Result<PathBuf, ()> {
    let data_local_dir =
        dirs::data_local_dir().ok_or_else(|| error!("failed to find data directory"))?;

    Ok(data_local_dir.join("tq"))
}

fn task_stdout_file(id: usize) -> Result<PathBuf, ()> {
    let output = data_dir()?.join("output");
    let _ = std::fs::create_dir_all(&output);

    Ok(output.join(format!("{}.stdout", id)))
}

fn task_stderr_file(id: usize) -> Result<PathBuf, ()> {
    let output = data_dir()?.join("output");
    let _ = std::fs::create_dir_all(&output);

    Ok(output.join(format!("{}.stderr", id)))
}

fn task_info_file(id: usize) -> Result<PathBuf, ()> {
    let info_dir = data_dir()?.join("info");
    let _ = std::fs::create_dir_all(&info_dir);

    Ok(info_dir.join(format!("{}.json", id)))
}

fn runtime() -> Result<(), ()> {
    let cli = Cli::parse();

    env_logger::init_from_env("TQ_LOG");

    match cli {
        Cli::StartSrv { force } => {
            let (tx, rx) = std::sync::mpsc::channel();
            std::thread::spawn(move || listen_message(tx, force));
            Srv::new(rx).run()?
        }
        Cli::StopSrv => send_message(&Message::Stop)?,

        Cli::Add { prog, args, name } => runtime_add(name, prog, args)?,

        Cli::Stat { id } => runtime_stat(id)?,
        Cli::Log { id, stderr } => runtime_log(stderr, id)?,
        Cli::Kill { id } => send_message(&Message::Kill { id })?,
    }

    Ok(())
}

fn runtime_add(show_name: Option<String>, prog: String, args: Vec<String>) -> Result<(), ()> {
    let cwd = std::env::current_dir().map_err(|err| error!("failed to get cwd: {err}"))?;
    let envs = std::env::vars().collect();

    let message = Message::Add(MessageAdd {
        show_name,
        envs,
        cwd,
        prog,
        args,
    });

    send_message(&message)?;
    Ok(())
}

fn runtime_log(stderr: bool, id: usize) -> Result<(), ()> {
    if stderr {
        let mut stderr_file = File::open(task_stderr_file(id)?)
            .map_err(|err| error!("failed to open stderr file: {err}"))?;

        std::io::copy(&mut stderr_file, &mut std::io::stdout())
            .map_err(|err| error!("failed to read task stderr: {err}"))?;
    } else {
        let mut stdout_file = File::open(task_stdout_file(id)?)
            .map_err(|err| error!("failed to open stdout file: {err}"))?;

        std::io::copy(&mut stdout_file, &mut std::io::stdout())
            .map_err(|err| error!("failed to read task stdout: {err}"))?;
    }

    Ok(())
}

fn runtime_stat(id: usize) -> Result<(), ()> {
    let rdr = File::open(task_info_file(id)?)
        .map_err(|err| error!("failed to open task info file: {err}"))?;

    let info: TaskInfo = serde_json::from_reader(rdr)
        .map_err(|err| error!("failed to parse task info file: {err}"))?;

    println!("TaskID: {:?}", info.id);
    println!("TaskStatus: {:?}", info.status);
    println!("Command: {} {}", info.prog, info.args.join(" "));
    println!("Cwd: {}", info.cwd.display());
    println!("Pid: {}", info.pid);

    Ok(())
}

fn send_message(message: &Message) -> Result<(), ()> {
    let mut stream = UnixStream::connect(SOCKET_FILE)
        .map_err(|err| error!("failed to connect to {SOCKET_FILE:?}: {err}"))?;

    stream
        .write_all(
            serde_json::to_string(message)
                .map_err(|err| error!("failed to serialize message: {err}"))?
                .as_bytes(),
        )
        .map_err(|err| error!("failed to write message: {err}"))?;

    Ok(())
}

const SOCKET_FILE: &str = "/tmp/tq.sock";

#[derive(serde::Deserialize, serde::Serialize)]
struct MessageAdd {
    show_name: Option<String>,
    prog: String,
    args: Vec<String>,
    cwd: PathBuf,
    envs: Vec<(String, String)>,
}

#[derive(serde::Deserialize, serde::Serialize)]
enum Message {
    Add(MessageAdd),
    Kill { id: usize },
    Stop,
}

struct Srv {
    tasks: Vec<Task>,
    rx: std::sync::mpsc::Receiver<Message>,
}

fn listen_message(tx: std::sync::mpsc::Sender<Message>, force: bool) -> Result<(), ()> {
    if force {
        let _ = std::fs::remove_file(SOCKET_FILE);
    }

    for client in UnixListener::bind(SOCKET_FILE)
        .map_err(|err| error!("failed to bind to {}: {}", SOCKET_FILE, err))?
        .incoming()
    {
        let mut client = client.map_err(|err| error!("failed to open client: {err}"))?;

        let mut data = String::new();
        client
            .read_to_string(&mut data)
            .map_err(|err| error!("failed to read from client: {err}"))?;

        let message: Message = serde_json::from_str(&data)
            .map_err(|err| error!("failed to parser message from client: {err}"))?;

        tx.send(message)
            .map_err(|err| error!("failed to send message to server: {err}"))?;
    }
    Ok(())
}

impl Srv {
    fn new(rx: std::sync::mpsc::Receiver<Message>) -> Self {
        Self { tasks: vec![], rx }
    }

    fn handle_message(&mut self, message: Message) -> Result<bool, ()> {
        match message {
            Message::Add(msg) => {
                if self.add_task(&msg).is_err() {
                    error!("failed to add task");
                    return Err(());
                }
            }
            Message::Stop => {
                if std::fs::remove_file(SOCKET_FILE).is_err() {
                    error!("failed to remove {SOCKET_FILE}");
                    return Err(());
                }
                return Ok(true);
            }
            Message::Kill { id } => {
                let Some(task) = self.tasks.get_mut(id) else {
                    error!("invalid task id: {id}");
                    return Err(());
                };

                if let Err(err) = task.child.kill() {
                    error!("can not kill task: {err}");
                    return Err(());
                }

                task.info.change_status(Status::Killed)?;
            }
        }

        Ok(false)
    }

    fn run(&mut self) -> Result<(), ()> {
        loop {
            match self.rx.recv_timeout(std::time::Duration::from_millis(100)) {
                Ok(message) => {
                    if self.handle_message(message)? {
                        break;
                    }
                }
                Err(std::sync::mpsc::RecvTimeoutError::Timeout) => (),
                Err(_) => {
                    error!("failed to recv message");
                    return Err(());
                }
            }

            for task in self.tasks.iter_mut() {
                if matches!(task.info.status, Status::Done(_) | Status::Killed) {
                    continue;
                }

                if let Some(status) = task
                    .child
                    .try_wait()
                    .map_err(|err| error!("try wait {} failed: {err}", task.info.prog))?
                {
                    let code = status
                        .code()
                        .ok_or_else(|| error!("failed to get return code of {}", task.info.prog))?;

                    task.info.change_status(Status::Done(code))?;
                };
            }

            self.task_menu();
        }

        Ok(())
    }

    fn task_menu(&mut self) {
        if self.tasks.is_empty() {
            print!("No Task.\r");
        } else {
            print!("Tasks:");

            for task in self.tasks.iter() {
                print!(" ");

                task_in_task_menu(task);
            }
            print!("\r");
        }

        let _ = std::io::stdout().flush();
    }

    fn add_task(&mut self, msg: &MessageAdd) -> Result<(), ()> {
        let id = self.tasks.len();

        let child = std::process::Command::new(&msg.prog)
            .args(&msg.args)
            .stdout(
                File::create(task_stdout_file(id)?)
                    .map_err(|err| error!("failed to open stdout file: {err}"))?,
            )
            .stderr(
                File::create(task_stderr_file(id)?)
                    .map_err(|err| error!("failed to open stderr file: {err}"))?,
            )
            .stdin(std::process::Stdio::null())
            .current_dir(&msg.cwd)
            .envs(msg.envs.to_owned())
            .spawn()
            .map_err(|err| error!("failed to spawn {prog:?}: {err}", prog = msg.cwd))?;

        let info = TaskInfo {
            pid: child.id(),
            show_name: msg.show_name.clone(),
            id,
            status: Status::Running,
            prog: msg.prog.to_owned(),
            args: msg.args.to_owned(),
            envs: msg.envs.to_owned(),
            cwd: msg.cwd.to_owned(),
        };

        dump_task_info(&info)?;

        self.tasks.push(Task { info, child });

        Ok(())
    }
}

fn task_in_task_menu(task: &Task) {
    match task.info.status {
        Status::Running => {
            print!(
                "{}:{}",
                task.info.id,
                task.info.show_name.as_ref().unwrap_or(&task.info.prog)
            );
        }
        Status::Killed => {
            let style = colorz::Style::new().fg(colorz::ansi::Yellow);

            print!(
                "{color}{}:{}{reset}",
                task.info.id,
                task.info.show_name.as_ref().unwrap_or(&task.info.prog),
                color = style.apply(),
                reset = style.clear()
            );
        }
        Status::Done(code) => {
            if code == 0 {
                let style = colorz::Style::new().fg(colorz::ansi::Green);

                print!(
                    "{color}{}:{}{reset}",
                    task.info.id,
                    task.info.prog,
                    color = style.apply(),
                    reset = style.clear()
                );
            } else {
                let style = colorz::Style::new().fg(colorz::ansi::Red);

                print!(
                    "{color}{}({code}):{}{reset}",
                    task.info.id,
                    task.info.prog,
                    color = style.apply(),
                    reset = style.clear()
                );
            }
        }
    }
}

fn dump_task_info(info: &TaskInfo) -> Result<(), ()> {
    let info_file_path = task_info_file(info.id)?;
    let info_file_path_tmp = info_file_path.with_extension("tmp");

    let info_file = File::create(&info_file_path_tmp)
        .map_err(|err| error!("failed to open info file: {err}"))?;

    serde_json::to_writer(info_file, info).map_err(|err| error!("failed to write info: {err}"))?;

    std::fs::rename(&info_file_path_tmp, &info_file_path)
        .map_err(|err| error!("failed to rename tmp info to info file: {err}"))?;
    Ok(())
}
